# -*- coding: utf-8 -*-

from __future__ import unicode_literals
from __future__ import print_function

from operator import add
from typing import Dict, List
from .util import merge_entity_attributes

from pyspark import RDD


def _merge_full_join_item(x):
    if x[0] is None:
        v = []
    else:
        v = x[0]
    return v


def pagerank(link_rdd, num_partitions, iterations=10):  # type: (RDD, int, int) -> RDD
    """
    Calculate pagerank for nodes in `links_rdd`
    
    links_rdd: [(in, out)]
    num_partitions: the default number of partitions to use
    iterations: iterations
    returns: an rdd of [(kgid, pagerank)]
    """
    if num_partitions is None:
        num_partitions = link_rdd.getNumPartitions()

    links = link_rdd.map(lambda x: (x["in"], x["out"]))

    links_in = links.filter(lambda x: x[0] != x[1]).distinct(num_partitions).groupByKey()
    links_out = links.map(lambda x: (x[1], None)).distinct(num_partitions)

    # we construct the full link matrix (instead of partial)
    links_collection = links_in.fullOuterJoin(links_out).mapValues(_merge_full_join_item).persist()
    
    # print("Total entities by links:", links_collection.count())
    ranks = links_collection.mapValues(lambda x: 1.0)

    for i in range(iterations):
        ranks = links_collection.join(ranks).flatMap(
            lambda x: [(d, x[1][1]/len(x[1][0])) for d in x[1][0]]
        ).reduceByKey(add, num_partitions).mapValues(
            lambda x: 0.15+0.85*x
        )
        ranks = links_collection.leftOuterJoin(ranks).mapValues(lambda x: x[1] if x[1] else 0.15)

    return ranks


def add_default_page_rank_score(x):
    if "pageRankScore" not in x:
        x["pageRankScore"] = 0.15
    return x


def add_default_entity_score(x):
    if "entityScore" not in x:
        x["entityScore"] = int(x.get("pageRankScore", 0.15) * 100)
    return x


def calc(entity, link, **kwargs):  # type: (RDD, RDD, Dict) -> List
    partition = kwargs.get("partition")
    iterations = kwargs.get("iterations", 5)
    default_only = kwargs.get("default_only", False)

    e = entity.map(lambda x: (x["@id"], x))

    if not default_only:
        if not link.isEmpty():
            pr = pagerank(link, partition, iterations)
            e = merge_entity_attributes(e, pr, "pageRankScore")

    e = e.mapValues(add_default_page_rank_score)
    e = e.mapValues(add_default_entity_score)
    return [e.map(lambda x: x[1]), link]
